// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

class EnforceSingleRowTest : public OperatorTestBase {
 protected:
  void assertQueryFails(
      const std::shared_ptr<const core::PlanNode>& plan,
      const std::string& errorMessage) {
    CursorParameters params;
    params.planNode = plan;
    try {
      readCursor(params, [](auto /*task*/) {});
      FAIL() << "Expected query to fail, but it succeeded";
    } catch (const PolluxException& e) {
      ASSERT_TRUE(e.message().find(errorMessage) != std::string::npos)
          << "Expected query to fail with error message: " << errorMessage
          << ". The query failed with a different error message: "
          << e.message();
    }
  }
};

TEST_F(EnforceSingleRowTest, basic) {
  auto singleRow =
      make_row_vector({make_flat_vector<int32_t>(1, [](auto row) { return row; })});
  auto multipleRows = make_row_vector(
      {make_flat_vector<int32_t>(27, [](auto row) { return row; })});

  createDuckDbTable({singleRow});

  // Single row of input. The query should pass.
  auto plan = PlanBuilder().values({singleRow}).enforceSingleRow().planNode();
  assertQuery(plan, "SELECT * FROM tmp");

  // Two rows of input in two separate single-row batches. The query should
  // fail.
  assertQueryFails(
      PlanBuilder()
          .values({singleRow, singleRow})
          .enforceSingleRow()
          .planNode(),
      "Expected single row of input. Received 1 extra rows.");

  // Multiple rows of input in a single batch. The query should fail.
  assertQueryFails(
      PlanBuilder().values({multipleRows}).enforceSingleRow().planNode(),
      "Expected single row of input. Received 27 rows.");

  // Empty input. The query should pass and return a single row of nulls.
  plan = PlanBuilder()
             .values({singleRow})
             .filter("c0 = 12345")
             .enforceSingleRow()
             .planNode();
  assertQuery(plan, "SELECT null");
}
